from config import config_singleton, config_model_singleton
import importlib
from threading import Thread
import concurrent.futures
import multiprocessing
from utils import sys_logger
# models

from models.base_model import ModelConfig, BaseModel
from models.constant_vals_admodel import ConstantVals_AnomalyDetector
from models.electric_servo_admodel import Electric_Servo_AnomalyDetector
from models.hydraulic_servo1d_admodel import Hydraulic_Servo1D_AnomalyDetector
from models.hydraulic_servo2d_admodel import Hydraulic_Servo2D_AnomalyDetector
from models.imu_admodel import IMU_AnomalyDetector


class ModelManager:
    """
    模型管理模块,具备管理,启动,停止模型线程的能力
    """

    def __init__(self, configs):

        self._manager = dict()

        self._model_configs_list = list()

        # 存储线程, 防止重复启动线程
        self._thread_dict = dict()  # 存储线程
        self.executor = None

        for config in configs:
            model_config = ModelConfig()
            model_config.name = config["name"]
            model_config.module_file = config["module_file"]
            model_config.module_name = config["module_name"]
            model_config.model_weigths_file = config.get("model_weights_file", "")

            inputs = config["inputs"]
            for input in inputs:
                model_config.inputs[input["table_num"]] = input["var_indexs"]
                model_config.input_gaps[input["table_num"]] = input["bus_gap"]

            model_config.inputs_need_preprocess = config.get("inputs_need_preprocess", False)
            model_config.mean = config.get("mean", [])
            model_config.std = config.get("std", [])

            outputs = config["outputs"]
            model_config.outputs["table_num"] = outputs["table_num"]
            model_config.outputs["var_indexs"] = outputs["var_indexs"]
            model_config.outputs["var_names"] = outputs["var_names"]

            model_config.input_output_queue_size = config.get("input_output_queue_size", 100)
            model_config.online_training_interval = config.get("online_training_interval", 0)
            model_config.online_training_mode = config.get("online_train_mode", False)

            model_config.start_conditions = config.get("start_conditions", [])
            model_config.stop_conditions = config.get("stop_conditions", [])

            self._model_configs_list.append(model_config)
            model_config = None

    def get_module(self, config: ModelConfig):
        """
        """
        # base = importlib.import_module("models.base_model")
        module = None
        try:
            module = importlib.import_module(config.module_file)
            module = getattr(module, config.module_name)
        except ImportError:
            ModuleNotFoundError(config.module_file)
        
        if module is None:
            raise ValueError(f"Module {config.module_name} not Found.")

        return module(config)


    def create_models(self):
        """
        """
        for config in self._model_configs_list:
            model = self.get_module(config)
            for table_num, var_indexs in config.inputs.items():
                if self._manager.get(table_num) is None:
                    self._manager[table_num] = [model]
                else:
                    self._manager[table_num].append(model)
            # 获取启动条件
            if config.start_conditions != []:
                for condition in config.start_conditions:
                    table_num = condition["table_num"]
                    if self._manager.get(table_num, None) is None:
                        self._manager[table_num] = [model]
                    else:
                        self._manager[table_num].append(model)

            # 获取停止条件
            if config.stop_conditions != []:
                for condition in config.stop_conditions:
                    table_num = condition["table_num"]
                    if self._manager.get(table_num, None) is None:
                        self._manager[table_num] = [model]
                    else:
                        self._manager[table_num].append(model)


    def get_model(self, table_num) -> list:
        """
        """
        try:
            if table_num in self._manager.keys():
                models = self._manager[table_num]
            else:
                return None
        except KeyError:
            return None
        return models


    def run_all_models(self):
        """
        """
        # 获取计算机核心数量, 预留两个核心给数据解码和数据发送
        # num_cores = multiprocessing.cpu_count() - 2
        # 创建线程池，大小为num_cores
        # 所有的任务都是while循环，所以必须保证worker的数量和模型的数量匹配，不然后面的模型得不到运行
        needworkers = len(self._model_configs_list)
        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=needworkers)

        for table_num, models in self._manager.items():
            for model in models:
                # 防止重复启动线程
                name = model.get_model_name()
                if name in self._thread_dict.keys():
                    # sys_logger.info(f"Model {name} is already running.")
                    continue
                else:
                    future = self.executor.submit(self.run_model, model)
                    self._thread_dict[name] = future


    def run_model(self, model):
        """
        Wrapper function to run the model in a separate thread, and log its status
        """
        try:
            sys_logger.info(f"Model {model.get_model_name()} started.")
            model.run()
            sys_logger.info(f"Model {model.get_model_name()} finished.")
        except Exception as e:
            sys_logger.error(f"Error running model {model.get_model_name()}: {e}")




    def stop_all_models(self):
        """
        cleans up all running models and shuts down the thread pool
        """
        for models in self._manager.values():
            for model in models:
                if isinstance(model, BaseModel):
                    model.stop()

        # map(lambda x: x.stop(), self._manager.values())
        if self.executor is not None:
            self.executor.shutdown(wait=True)
            self._thread_dict.clear()
            sys_logger.info("All models have been stopped and thread pool has been shut down.")


    def get_all_models(self):

        return list(self._manager.values())




if __name__ == "__main__":

    modelconfigs = config_model_singleton.model_configs

    man = ModelManager(modelconfigs)
    man.create_models()

    print()
